#include "secure_rw.h"
#include <limits.h>
#include <string.h>
#include <openssl/aes.h>
#include <openssl/hmac.h>

#include <sys/types.h>

#include <sys/socket.h>
#include <sys/stat.h>
#include <fcntl.h>

static unsigned char *calculate_HMAC(unsigned int seq, const void *msg,
		size_t count, const unsigned char *secret);
static enum bool verify_HMAC(unsigned char *a, unsigned char *b);

ssize_t sc_write(struct connection *conn, void *buf, size_t count)
{
	if (conn->seq == UINT_MAX)
	{
		close(conn->sockfd);
		conn->connected = FALSE;
		return 0;
	}
	else if (count == 0)
		return 0;
	else if (conn->sockfd < 0)
		return -1;
	else
	{
		fprintf(stderr, "Seq #: %d\n", conn->seq);

		fprintf(stderr, "Message: ");
		print_hex(buf, count);
		fprintf(stderr, "\n");

		/* Init AES encryption key */
		AES_KEY enKey;
		AES_set_encrypt_key(conn->secret, SECRET_BITS, &enKey);

		/* Copy iv */
		unsigned char iv[CBC_BLOCK_SIZE];
		memcpy(iv, conn->iv, CBC_BLOCK_SIZE);

		/* Padding & Allocate buffer for encrypted text */
		size_t num_blocks = (count - 1) / 16 + 1;
		unsigned char last_block_size = count % CBC_BLOCK_SIZE;
		if (last_block_size == 0)
			num_blocks++;
		size_t encrypted_size = num_blocks * CBC_BLOCK_SIZE;
		unsigned char *padded = (unsigned char *) malloc(encrypted_size);
		memcpy(padded, buf, count);
		*(padded + encrypted_size - 1) = last_block_size;
		unsigned char *encrypted = (unsigned char *) malloc(encrypted_size);

		/* Encrypt */
		AES_cbc_encrypt(padded, encrypted, encrypted_size, &enKey, iv,
				AES_ENCRYPT);
		fprintf(stderr, "Encrypted: ");
		print_hex(encrypted, encrypted_size);
		fprintf(stderr, "\n");

		unsigned char *hmac = calculate_HMAC(conn->seq, buf, count,
				conn->secret);

		fprintf(stderr, "HMAC: ");
		print_hex(hmac, MD_SIZE);
		fprintf(stderr, "\n");

		/* Prepare (i || c || a) */
		size_t ica_size = sizeof(conn->seq) + encrypted_size + MD_SIZE;
		unsigned char *ica = (unsigned char *) malloc(ica_size);
		memcpy(ica, &conn->seq, sizeof(conn->seq));
		memcpy(ica + sizeof(conn->seq), encrypted, encrypted_size);
		memcpy(ica + sizeof(conn->seq) + encrypted_size, hmac, MD_SIZE);

		/* Write to physical channel */
//		printf("Size: %d Final: ", ica_size);
//		print_hex(ica, ica_size);
//		printf("\n");
		ssize_t r = send(conn->sockfd, ica, ica_size, MSG_NOSIGNAL);
		free(encrypted);
		free(ica);
		free(padded);

		/* Return */
		if (r == ica_size)
		{
			conn->seq++;
			return count;
		}
		else
			return -1;
	}
}

ssize_t sc_read(struct connection *conn, void *buf, size_t count)
{
	unsigned char tmp[MAX_READ];
	ssize_t len = read(conn->sockfd, tmp, MAX_READ);
//	printf("Read: ");
//	print_hex(tmp, len);
//	printf("\n");

	/* Check sequence number */
	unsigned int seq = *(unsigned int *) tmp;
	fprintf(stderr, "Seq #: %d, Expected #: %d\n", seq, conn->seq);
	if (seq == conn->seq)
	{
		conn->seq++;
	}
	else
		return -1;

	/* Get c and a */
	unsigned char *encrypted = tmp + sizeof(conn->seq);
	unsigned char *md = tmp + len - MD_SIZE;
	fprintf(stderr, "Encrypted: ");
	print_hex(encrypted, md - encrypted);
	fprintf(stderr, "\n");

	/* Init AES decryption key */
	AES_KEY deKey;
	AES_set_decrypt_key(conn->secret, SECRET_BITS, &deKey);

	/* Allocate buffer for decrypted text */
	size_t decrypted_size = md - encrypted;
	unsigned char *decrypted = (unsigned char *) malloc(decrypted_size);

	/* Copy iv */
	unsigned char iv[CBC_BLOCK_SIZE];
	memcpy(iv, conn->iv, CBC_BLOCK_SIZE);

	/* Decrypt */
	AES_cbc_encrypt(encrypted, decrypted, decrypted_size, &deKey, iv,
			AES_DECRYPT);
//	printf("Decrypted: ");
//	print_hex(decrypted, decrypted_size);
//	printf("\n");

	/* Unpad */
	unsigned char last_block_size = *(decrypted + decrypted_size - 1);
	size_t msglen = decrypted_size - (CBC_BLOCK_SIZE - last_block_size);
	fprintf(stderr, "Message: ");
	print_hex(decrypted, msglen);
	fprintf(stderr, "\n");

	fprintf(stderr, "HMAC: ");
	print_hex(md, MD_SIZE);
	fprintf(stderr, "\n");

	unsigned char *hmac = calculate_HMAC(seq, decrypted, msglen, conn->secret);
	if (verify_HMAC(hmac, md))
	{
		fprintf(stderr, "HMAC verified!\n");
		memcpy(buf, decrypted, msglen);
		free(decrypted);
		return msglen;
	}
	else
	{
		fprintf(stderr, "HMAC error!\n");
		free(decrypted);
		return -1;
	}
}

static unsigned char *calculate_HMAC(unsigned int seq, const void *msg,
		size_t count, const unsigned char *secret)
{
	/* Prepare (i || m) */
	unsigned char *im = (unsigned char *) malloc(sizeof(seq) + count);
	memcpy(im, &seq, sizeof(seq));
	memcpy(im + sizeof(seq), msg, count);

	/* Calculate HMAC */
	unsigned char *hmac;
	hmac = HMAC(EVP_sha256(), secret, SECRET_BITS / 8, im, sizeof(seq) + count,
			NULL, NULL);
	free(im);
	return hmac;
}

static enum bool verify_HMAC(unsigned char *a, unsigned char *b)
{
	int i;
	for (i = 0; i < MD_SIZE; ++i)
	{
		if (a[i] != b[i])
			return FALSE;
	}
	return TRUE;
}

void print_hex(unsigned char *hex, size_t len)
{
	size_t i;
	for (i = 0; i < len; i++)
		fprintf(stderr, "%02x", hex[i]);
}
